import { axisColorConfig } from '@/config/stage-canvas-config';
import { ThreeEvent, useThree } from '@react-three/fiber';
import { omit } from 'lodash-es';
import * as React from 'react';
import { AxesHelper, CanvasTexture, Color, Vector3 } from 'three';

type AxisProps = {
  color: string;
  rotation: [number, number, number];
  scale?: number;
};

type AxisHeadProps = JSX.IntrinsicElements['sprite'] & {
  arcStyle: string;
  label?: string;
  labelColor: string;
  axisHeadScale?: number;
  disabled?: boolean;
  font: string;
  onClick?: (e: ThreeEvent<MouseEvent>) => null;
};

type GizmoViewportProps = JSX.IntrinsicElements['group'] & {
  axisColors?: [string, string, string];
  axisScale?: number;
  labels?: [string, string, string];
  axisHeadScale?: number;
  labelColor?: string;
  hideNegativeAxes?: boolean;
  hideAxisHeads?: boolean;
  disabled?: boolean;
  font?: string;
  onClick?: (e: ThreeEvent<MouseEvent>) => null;
};

function Axis({ scale = 1, color, rotation }: AxisProps) {
  const dir = new Vector3(1, 0, 0);
  const origin = new Vector3(0, 0, 0);
  const targetLength = 1 * scale;
  const headLength = 0.2;
  const headWidth = headLength * 0.5;
  const length = targetLength + headLength;

  return (
    <group rotation={rotation}>
      <arrowHelper args={[dir, origin, length, color, headLength, headWidth]} />
    </group>
  );
}

function AxisHead({
  onClick,
  font,
  disabled,
  arcStyle,
  label,
  labelColor,
  axisHeadScale = 1,
  ...props
}: AxisHeadProps) {
  const gl = useThree((state) => state.gl);
  const texture = React.useMemo(() => {
    const canvas = document.createElement('canvas');
    canvas.width = 64;
    canvas.height = 64;

    const context = canvas.getContext('2d')!;

    if (label) {
      context.font = font;
      context.textAlign = 'center';
      context.fillStyle = labelColor;
      context.fillText(label, 32, 41);
    }
    return new CanvasTexture(canvas);
  }, [arcStyle, label, labelColor, font]);

  const scale = (label ? 1 : 0.75) * 1 * axisHeadScale;

  return (
    <sprite scale={scale} {...omit(props, 'onClick')}>
      <spriteMaterial
        map={texture}
        map-anisotropy={gl.capabilities.getMaxAnisotropy() || 1}
        alphaTest={0.3}
        opacity={label ? 1 : 0.75}
        toneMapped={false}
      />
    </sprite>
  );
}

export const GizmoViewport = ({
  hideNegativeAxes,
  hideAxisHeads,
  disabled,
  font = '18px Inter var, Arial, sans-serif',
  axisColors = [axisColorConfig.x, axisColorConfig.y, axisColorConfig.z],
  axisHeadScale = 1,
  axisScale,
  labels = ['X', 'Y', 'Z'],
  labelColor = '#000',
  onClick,
  ...props
}: GizmoViewportProps) => {
  const [colorX, colorY, colorZ] = axisColors;
  const axisHeadProps = {
    font,
    disabled,
    labelColor,
    onClick,
    axisHeadScale,
  };

  const axesHelperRef = React.useRef<AxesHelper>(null);
  React.useEffect(() => {
    const axesHelper = axesHelperRef.current;
    axesHelper?.setColors(
      new Color(colorX),
      new Color(colorY),
      new Color(colorZ),
    );
  }, []);
  return (
    <group scale={40} {...props}>
      <Axis color={colorX} rotation={[0, 0, 0]} scale={axisScale} />
      <Axis color={colorY} rotation={[0, 0, Math.PI / 2]} scale={axisScale} />
      <Axis color={colorZ} rotation={[0, -Math.PI / 2, 0]} scale={axisScale} />
      {!hideAxisHeads && (
        <>
          <AxisHead
            arcStyle={colorX}
            position={[1.5, 0, 0]}
            label={labels[0]}
            {...axisHeadProps}
            labelColor={axisColorConfig.x}
          />
          <AxisHead
            arcStyle={colorY}
            position={[0, 1.5, 0]}
            label={labels[1]}
            {...axisHeadProps}
            labelColor={axisColorConfig.y}
          />
          <AxisHead
            arcStyle={colorZ}
            position={[0, 0, 1.5]}
            label={labels[2]}
            {...axisHeadProps}
            labelColor={axisColorConfig.z}
          />
        </>
      )}
    </group>
  );
};
